from abc import ABC
from enum import Enum, auto
from typing import Union, List, Dict

# from vm.chunk import Chunk


class Value:
    def __init__(self, value_type: 'ValueType', value: Union[None, str, float, 'Obj', 'Value'] = None):
        self.value_type = value_type
        self.value = value

    def __bool__(self):
        match self.value_type:
            case ValueType.BOOL:
                return self.value
            case ValueType.NIL:
                return False
            case ValueType.NUMBER:
                return self.value == 0
            case ValueType.OBJ:
                if self.value.obj_type == ObjType.UPVALUE:
                    return bool(self.value.value)
                else:
                    return True

    def __repr__(self):
        match self.value_type:
            case ValueType.NIL:
                return 'Nil'
            case ValueType.NUMBER:
                return f'{self.number_value()}'
            case ValueType.BOOL:
                return f'{str(bool(self))}'
            case ValueType.OBJ:
                return repr(self.value)
            case _:
                raise TypeError(f'invalid value type {self.value_type}')

    @classmethod
    def nil(cls):
        return cls(ValueType.NIL, None)

    @classmethod
    def new_number(cls, n: float):
        return cls(ValueType.NUMBER, n)

    @classmethod
    def boolean(cls, b: bool):
        return cls(ValueType.BOOL, b)

    @classmethod
    def wrap(cls, v: 'Value'):
        return cls(ValueType.OBJ, Obj.upvalue(v))

    def unwrap(self):
        return self.value.unwrap()

    def update_upvalue(self, v: 'Value'):
        self.value.update_upvalue(v)

    @classmethod
    def string(cls, s: str):
        return cls(ValueType.OBJ, Obj.string(s))

    @classmethod
    def klass(cls, name: str):
        return cls(ValueType.OBJ, Obj.klass(name))

    @classmethod
    def closure(cls, fn: 'ObjFunc', upvalues: List['Value']):
        return cls(ValueType.OBJ, Obj.closure(fn, upvalues))

    def number_value(self) -> float:
        if self.value_type == ValueType.NUMBER:
            return self.value
        raise TypeError()

    def string_value(self) -> str:
        if self.value_type == ValueType.OBJ:
            return self.value.string_value()
        raise TypeError()

    def func_value(self) -> 'ObjFunc':
        if self.value_type == ValueType.OBJ:
            return self.value.func_value()

        raise TypeError(f'invalid function value, actually {self.value_type}')

    def closure_value(self) -> 'ObjClosure':
        if self.value_type == ValueType.OBJ:
            return self.value.closure_value()

        raise TypeError(f'invalid closure value, actually {self.value_type}')

    def class_value(self) -> 'ObjClass':
        if self.value_type == ValueType.OBJ:
            return self.value.class_value()

        raise TypeError(f'invalid class value, actually {self.value_type}')

    def instance_value(self) -> 'ObjInstance':
        if self.value_type == ValueType.OBJ:
            return self.value.instance_value()

        raise TypeError(f'invalid instance value, actually {self.value_type}')


class ValueType(Enum):
    BOOL = auto()
    NIL = auto()
    NUMBER = auto()
    OBJ = auto()


class Obj(ABC):
    def __init__(self, obj_type: 'ObjType'):
        self.obj_type = obj_type
        self.marked = False

    @classmethod
    def string(cls, s: str):
        return ObjString(s)

    @classmethod
    def upvalue(cls, value: 'Value'):
        return ObjUpvalue(value)

    @classmethod
    def func(cls, name: str, arity: int, upvalue_count: int, chunk: 'Chunk'):
        return ObjFunc(name, arity, upvalue_count, chunk)

    @classmethod
    def closure(cls, f: 'ObjFunc', upvalues: List[Value]):
        return ObjClosure(f, upvalues)

    @classmethod
    def klass(cls, name: str):
        return ObjClass(name)

    @classmethod
    def instance(cls, klass: 'ObjClass'):
        return ObjInstance(klass)

    @classmethod
    def bound_method(cls, receiver: 'Value', method: 'ObjClosure'):
        return ObjBoundMethod(receiver, method)

    def mark(self, marked: bool):
        self.marked = marked

    def string_value(self) -> str:
        if self.obj_type != ObjType.STRING:
            raise TypeError(f'cannot get string value for {self.obj_type}')
        return self

    def func_value(self) -> 'ObjFunc':
        if self.obj_type != ObjType.FUNCTION:
            raise TypeError(f'cannot get function value for {self.obj_type}')
        return self

    def closure_value(self) -> 'ObjClosure':
        if self.obj_type != ObjType.CLOSURE:
            raise TypeError(f'cannot get closure value for {self.obj_type}')
        return self

    def upvalue_value(self) -> 'ObjUpvalue':
        if self.obj_type != ObjType.UPVALUE:
            raise TypeError(f'cannot update upvalue for {self.obj_type}')
        return self

    def class_value(self) -> 'ObjClass':
        if self.obj_type != ObjType.CLASS:
            raise TypeError(f'cannot get class value for {self.obj_type}')
        return self

    def instance_value(self) -> 'ObjInstance':
        if self.obj_type != ObjType.INSTANCE:
            raise TypeError(f'cannot get instance value for {self.obj_type}')
        return self

    def bound_method_value(self) -> 'ObjBoundMethod':
        if self.obj_type != ObjType.BOUND_METHOD:
            raise TypeError(f'cannot get bound method value for {self.obj_type}')
        return self


class ObjString(Obj):
    def __init__(self, s: str):
        super(ObjString, self).__init__(ObjType.FUNCTION)
        self.value = s

    def string_value(self) -> str:
        return self.value

    def __repr__(self):
        return self.value


class ObjClass(Obj):
    def __init__(self, name: str):
        super(ObjClass, self).__init__(ObjType.CLASS)
        self.name = name
        self.methods: Dict[str, 'ObjClosure'] = {}

    def add_method(self, name: str, value: 'ObjClosure'):
        self.methods[name] = value

    def mark(self, marked: bool):
        super(ObjClass, self).mark(marked)
        for closure in self.methods.values():
            closure.mark(marked)

    def class_value(self) -> 'ObjClass':
        return self

    def __repr__(self):
        return f'Class(name={self.name}, methods={list(self.methods.keys())})'


class ObjFunc(Obj):
    def __init__(self, name: str, arity: int, upvalue_count: int, chunk: 'Chunk'):
        super(ObjFunc, self).__init__(ObjType.FUNCTION)
        self.name = name
        self.arity = arity
        self.upvalue_count = upvalue_count
        self.chunk = chunk

    def func_value(self) -> 'ObjFunc':
        return self

    def __repr__(self):
        return f'ObjFunc(name={self.name})'


class ObjClosure(Obj):
    def __init__(self, fun: 'ObjFunc', upvalues: List[Value]):
        super(ObjClosure, self).__init__(ObjType.CLOSURE)
        self.fun = fun
        self.upvalues = upvalues

    def mark(self, marked: bool):
        super(ObjClosure, self).mark(marked)
        for value in self.upvalues:
            if value.value_type == ValueType.OBJ:
                value.value.mark(marked)

    def closure_value(self) -> 'ObjClosure':
        return self

    def __repr__(self):
        return f'ObjClosure(name={self.fun.name}, upvalues={self.upvalues})'


class ObjInstance(Obj):
    def __init__(self, klass: 'ObjClass'):
        super(ObjInstance, self).__init__(ObjType.INSTANCE)
        self.klass = klass
        self.fields: Dict[str, 'Value'] = {}

    def mark(self, marked: bool):
        super(ObjInstance, self).mark(marked)
        self.klass.mark(marked)
        for value in self.fields.values():
            if value.value_type == ValueType.OBJ:
                value.value.mark(marked)

    def instance_value(self) -> 'ObjInstance':
        return self

    def __repr__(self):
        return f'ObjInstance(class={self.klass.name}, fields={self.fields})'


class ObjBoundMethod(Obj):
    def __init__(self, receiver: 'Value', method: 'ObjClosure'):
        super(ObjBoundMethod, self).__init__(ObjType.BOUND_METHOD)
        self.receiver = receiver
        self.method = method

    def mark(self, marked: bool):
        super(ObjBoundMethod, self).mark(marked)
        if self.receiver.value_type == ValueType.OBJ:
            self.receiver.value.mark(marked)

    def __repr__(self):
        return f'ObjBoundMethod(receiver={self.receiver}, method={self.method.fun.name})'


class ObjUpvalue(Obj):
    def __init__(self, value: 'Value'):
        super(ObjUpvalue, self).__init__(ObjType.UPVALUE)
        self.value = value

    def mark(self, marked: bool):
        super(ObjUpvalue, self).mark(marked)
        if self.value.value_type == ValueType.OBJ:
            self.value.value.mark(marked)

    def update_upvalue(self, value: 'Value'):
        self.value = value

    def unwrap(self) -> 'Value':
        return self.value


class ObjType(Enum):
    STRING = auto()
    CLASS = auto()
    BOUND_METHOD = auto()
    FUNCTION = auto()
    CLOSURE = auto()
    INSTANCE = auto()
    NATIVE = auto()
    UPVALUE = auto()


